#!/usr/bin/perl -w
# AUTHOR: Divya Sain <divyasain-AT-gmail.com>
#         Jason Stajich <jason.stajich-AT-ucr.edu>

use Getopt::Long;

my $all = 0;
GetOptions(
	'a|all!' => \$all,
	);
# provide the mclresults file as an argument or defaults to 'mclresults' as the name
my $filename = shift || 'mclresults';
open(IN, $filename) || die "Hey! I cannot open $filename: '$!'\n";
my $OGnum = 1;
my (%gene2OG,%OG2gene,%sp_prefix);

# parse the ortholog group file
# when this is done we can lookup a gene and see its ortholog group
# in %gene2OG
# and we can find all the genes in a particular ortholog group in 
# %OG2gene 

while(<IN>) {
    my (@genes) = split;
    for my $gene ( @genes )  {
	$gene2OG{$gene} = $OGnum;
	my ($sp) = substr($gene,0,2);
	$sp_prefix{$sp}++;
    }
    @{$OG2gene{$OGnum}} = @genes;
    $OGnum++;
}

#while (($gene, $OGnum) = each %gene2OG){
	#print "$gene-> $OGnum\n";}	
#for my $OGnum ( sort { $a <=> $b } keys %OG2gene ) {
 #print "$OGnum: ", join(",", @{$OG2gene{$OGnum}}), "\n";}

# now let's identify genes that are differntially expressed in each 
# comparison
my @files = qw(nc_3hvs20h_filter.diff ro_3hvs20h_filter.diff cc_9dvs15d_filter.diff);
my %signif_genes;
for my $file ( @files ) {
    open( my $fh => $file) || die "Cannot open file '$file': $!\n";
    while(<$fh>) {
	next if /^test_id/;
	my ($genename,$foldchange,$status,$timepoint1,$timepoint2,$pvalue,$significant) = split;

	next unless $significant eq 'yes'; # only take the significantly differentially expressed ones
	# you may want to further refine this by taking only ones with P-value < certain amount like this
	# next unless $pvalue < 1e-3; 

	# store the gene which is significant in a hash, and track the fold change and pvalue
	$signif_genes{$genename} = [$foldchange,$pvalue];
    }
}

# now we can determine which genes are significant and which genes are in the same orthogroup

# which significant genes are in the same orthogroup?
my %keep;
for my $gene ( keys %signif_genes ) {
    my $orthogroup = $gene2OG{$gene};
    my @OGgenes = @{$OG2gene{$orthogroup}};
    
    for my $og_gene ( @OGgenes ) {
	if(exists $signif_genes{$og_gene} ) { # this gene is also significant
	    $keep{$orthogroup}->{$og_gene}++; # keep unique hash of all genes that are signif in an OG
	}
    }
    # also store our original gene, though this should be filled in anyways
    # later when we find it via one of the other genes in group which is significant
    # if we comment this out, we won't see any singleton OG
    $keep{$orthogroup}->{$gene}++;
}

my @spname = sort keys %sp_prefix;

print join("\t", 'OG', map { $_ ."_genes"} @spname),"\n";
for my $OG ( sort { $a <=> $b } keys %keep ) {
    # let's group them by their prefix
    my %byspecies;
    for my $gene ( sort keys %{$keep{$OG}} ) {
	my ($prefix) = substr($gene,0,2);
	push @{$byspecies{$prefix}}, $gene;
    }
    # check that there are 3 species, if there aren't skip it
    if( ! $all ) {
	next if (scalar keys %byspecies != 3);
    }
    print join("\t", $OG, map { join(",",@{$byspecies{$_} || []}) } @spname), "\n";
}
